import pandas as pd
from sklearn.cluster import k_means
import numpy as np
from scipy.spatial.distance import pdist, euclidean
from sklearn.datasets import load_iris


def DaviesBouldin(X, labels):
    n_cluster = len(np.bincount(labels))
    cluster_k = [X[labels == k] for k in range(n_cluster)]
    centroids = [np.mean(k, axis=0) for k in cluster_k]

    # 求S
    S = [np.mean([euclidean(p, centroids[i]) for p in k]) for i, k in enumerate(cluster_k)]
    Ri = []

    for i in range(n_cluster):
        Rij = []
        # 计算Rij
        for j in range(n_cluster):
            if j != i:
                r = (S[i] + S[j]) / euclidean(centroids[i], centroids[j])
                Rij.append(r)
        # 求Ri
        Ri.append(max(Rij))

        # 求dbi
    dbi = np.mean(Ri)

    return dbi


if __name__ == "__main__":
    data = load_iris()
    # print(data)
    print(data["data"].shape)
    print(data["target"].shape)
    print(data["data"])
    print(data["target"])
    res = DaviesBouldin(data["data"], data["target"])
    print(res)
